package com.ai.main.controller;

import jakarta.servlet.http.HttpServletResponse;
import lombok.RequiredArgsConstructor;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.vectorstore.QuestionAnswerAdvisor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.prompt.PromptTemplate;
import org.springframework.ai.template.st.StTemplateRenderer;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.http.MediaType;
import org.springframework.http.ResponseEntity;
import org.springframework.http.codec.ServerSentEvent;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Flux;

/**
 * 功能描述
 *
 * @author: jichao.huang
 * @date: 2025年04月22日 14:26
 */
@RequiredArgsConstructor
@RequestMapping("/api/chatClient")
@RestController
public class ChatController {
    private final ChatClient chatClient;

    private final VectorStore vectorStore;

    private static final String DEFAULT_PROMPT = "你好，介绍下你自己！";

    /**
    public ChatController(ChatClient.Builder builder) {
        this.chatClient = builder.build();
    }
     **/

    /**
    public ChatController(ChatModel chatModel) {
        chatClient = ChatClient.builder(chatModel)
                .defaultAdvisors(new SimpleLoggerAdvisor())
                .defaultOptions(OpenAiChatOptions.builder().topP(0.7).build())
                .build();

    }
     **/



    @GetMapping("/chat")
    public ResponseEntity<String> chat(@RequestParam("input") String input) {
        PromptTemplate customPromptTemplate = PromptTemplate.builder()
                .renderer(StTemplateRenderer.builder().startDelimiterToken('<').endDelimiterToken('>').build())
                .template("""
            <query>

            Context information is below.

			---------------------
			<question_answer_context>
			---------------------

			Given the context information and no prior knowledge, answer the query.

			Follow these rules:

			1. If the answer is not in the context, just say that you don't know.
			2. Avoid statements like "Based on the context..." or "The provided information...".
            """)
                .build();

        QuestionAnswerAdvisor qaAdvisor = QuestionAnswerAdvisor.builder(vectorStore)
                .promptTemplate(customPromptTemplate)
                .build();

        String conversationId = "789";
        try {
            ChatResponse response = this.chatClient.prompt()
                    .user(input)
                    //.advisors(a -> a.param(ChatMemory.CONVERSATION_ID, conversationId))
                    //.advisors(qaAdvisor)
                    .call()
                    .chatResponse();
            return ResponseEntity.ok(response.getResult().getOutput().getText());
        } catch (Exception e) {
            return ResponseEntity.internalServerError()
                    .body("AI服务请求失败: " + e.getMessage());
        }
    }

    @GetMapping("/simple/chat")
    public String simpleChat() {
        return this.chatClient.prompt(DEFAULT_PROMPT).call().content();
    }

    @GetMapping("/stream/chat")
    public Flux<String> streamChat(HttpServletResponse response) {
        response.setCharacterEncoding("utf-8");
        return this.chatClient.prompt(DEFAULT_PROMPT).stream().content();
    }

    @GetMapping(value = "/stream/chat2", produces = "text/html;charset=utf-8")
    public Flux<String> streamChat2(HttpServletResponse response) {
        response.setCharacterEncoding("utf-8");
        return this.chatClient.prompt(DEFAULT_PROMPT).stream().content();
    }

    @GetMapping(value = "/stream/chatMemory", produces = "text/html;charset=utf-8")
    public Flux<String> chatMemory(String prompt, String chatId) {
        return this.chatClient
                .prompt(prompt)
                .advisors(a -> a.param(ChatMemory.CONVERSATION_ID, chatId))
                .stream()
                .content();
    }

    @GetMapping(value = "/stream/rag", produces = "text/html;charset=utf-8")
    public Flux<String> chatRag(String prompt, String chatId) {
        return this.chatClient
                .prompt(prompt)
                .advisors(a -> a.param(QuestionAnswerAdvisor.FILTER_EXPRESSION, "type='Spring'"))
                .stream()
                .content();
    }

    @GetMapping(value = "/stream/response", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<ServerSentEvent<String>> streamResponse(@RequestParam("input") String input) {
        System.out.println("input:" + input);
        return this.chatClient.prompt().user(input).stream().content().map(content -> ServerSentEvent.<String>builder().data(content).build());
    }
}
